

import os
import sys
import shutil
import pandas as pd
from tqdm import tqdm
from moviepy.editor import AudioFileClip

sys.path.insert(0, "./asteroid/")
sys.path.append('./src/separationhandler/asteroid/')

from asteroid.models import BaseModel
from asteroid.models import ConvTasNet

from .tools.matcher import MatchHandler
from .tools.getfacelabel import get_face_label
from .tools.makeseplabel import make_separation_label


class Task3Handler:
    '''
    Args:
        @model_name:
            name of pretrained Separation Model
        @video_path:
            raw video of Task3
        @audio_path:
            convert video -> audio
        @face_label_path:
            get face ID using Face Recognition
        @match_table_path:
            a list of position(left, middle, right) and audio ID
        @sep_list_path:
            a list of old name(1, 2, 3) and new name(left, middle, right)
        @temp_audio_path:
            sep audio with name of _1, _2, _3
        @new_temo_audio_path:
            sep audio with name of _left, _middle, _right
    '''


    def __init__(self, model_name, video_path, audio_path, face_label_path, match_table_path, sep_list_path, temp_audio_path, new_temp_audio_path):
        
        self.model_name = model_name
        self.video_path = video_path
        self.audio_path = audio_path
        self.model = None
        self.face_label_path = face_label_path
        self.match_table_path = match_table_path
        self.sep_list_path = sep_list_path
        self.match_list = []
        self.temp_audio_path = temp_audio_path
        self.new_temp_audio_path = new_temp_audio_path


    def check_path(self):
        '''
            check path, if not exists, make it
        '''
        if not os.path.exists(self.audio_path):
            os.mkdir(self.audio_path)
        if not os.path.exists(self.temp_audio_path):
            os.mkdir(self.temp_audio_path)
        if not os.path.exists(self.temp_audio_path):
            os.mkdir(self.temp_audio_path)
        if not os.path.exists(self.new_temp_audio_path):
            os.mkdir(self.new_temp_audio_path)
    

    def process_video(self):
        '''
            video_path -> audio_path
        '''
        video_list = os.listdir(self.video_path)
        for video in video_list:
            video_path = os.path.join(self.video_path, video)
            audio_path = os.path.join(self.audio_path, video[:-4]+'.wav')
            if os.path.exists(audio_path):
                print('  Audio Path: (' + self.audio_path + ') Already Exists.')
                break
            curr_audio = AudioFileClip(video_path)
            curr_audio.write_audiofile(audio_path, logger=None)
            curr_audio.close()


    def load_model(self):
        '''
            get pretrained model
        '''
        self.model = ConvTasNet.from_pretrained(self.model_name)
        print('  Load Model: (' + self.model_name + ') Successful!')


    def separate_3_audio(self):
        '''
            audio_path -> temp_audio_path, separate the mixed audio
        '''
        mix_audio_list = os.listdir(self.audio_path)

        for mix_audio_idx in tqdm(range(len(mix_audio_list)), ncols=70):
            
            mix_audio = mix_audio_list[mix_audio_idx]
            mix_audio_path = os.path.join(self.audio_path, mix_audio)
            self.model.separate(mix_audio_path, resample=True, force_overwrite=True, output_dir=self.temp_audio_path)


    def rename_audio(self):
        '''
            process the raw name of audio by separator
        '''
        audio_list = os.listdir(self.temp_audio_path)
        for audio_idx in tqdm(range(len(audio_list)), ncols=70):
            audio = audio_list[audio_idx]
            audio_path = os.path.join(self.temp_audio_path, audio)
            new_path = audio_path.replace('est1', '1').replace('est2', '2').replace('est3', '3').replace('combine', '')
            os.rename(audio_path, new_path)


    def get_match_table(self):
        '''
            make table of (_1, _2, _3) and (_left, _middle, _right)
        '''
        label = get_face_label(self.face_label_path)
        handler = MatchHandler(self.temp_audio_path, label, self.match_table_path)
        for id in range(len(label)):
            self.match_list.append(handler.do_match_single(id+1))
        make_separation_label(self.match_list, self.sep_list_path, self.temp_audio_path, self.new_temp_audio_path)


    def add_audio_position(self):
        '''
            rename audio : (_1, _2, _3) -> (_left, _middle, _right)
        '''
        test_csv = pd.read_csv(self.sep_list_path, sep=',').values

        for noise_utterance, emb_utterance_id, clean_utterance, clean_utterance2 in test_csv:
            
            shutil.copy(noise_utterance, clean_utterance)


# # debug
# if __name__ == '__main__':

#     handler = Task3Handler(1, 2, 3, 4, 5)
#     # handler.get_match_table()
#     # os.system('python fursep.py')
#     handler.rename_audios()